// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <chrono>
#include <cstdint>

#include <nebula/compute/api_scalar.h>
#include <nebula/datetime/datetime.h>

namespace nebula::compute::internal {

    using alkaid_vendored::date::days;
    using alkaid_vendored::date::floor;
    using alkaid_vendored::date::local_days;
    using alkaid_vendored::date::local_time;
    using alkaid_vendored::date::locate_zone;
    using alkaid_vendored::date::sys_days;
    using alkaid_vendored::date::sys_time;
    using alkaid_vendored::date::time_zone;
    using alkaid_vendored::date::year_month_day;
    using alkaid_vendored::date::zoned_time;
    using std::chrono::duration_cast;

    inline int64_t GetQuarter(const year_month_day &ymd) {
        return static_cast<int64_t>((static_cast<uint32_t>(ymd.month()) - 1) / 3);
    }

    static inline turbo::Result<const time_zone *> LocateZone(const std::string &timezone) {
        try {
            return locate_zone(timezone);
        } catch (const std::runtime_error &ex) {
            return turbo::invalid_argument_error("Cannot locate timezone '", timezone, "': ", ex.what());
        }
    }

    static inline const std::string &GetInputTimezone(const DataType &type) {
        static const std::string no_timezone = "";
        switch (type.id()) {
            case Type::TIMESTAMP:
                return turbo::checked_cast<const TimestampType &>(type).timezone();
            default:
                return no_timezone;
        }
    }

    static inline turbo::Status ValidateDayOfWeekOptions(const DayOfWeekOptions &options) {
        if (options.week_start < 1 || 7 < options.week_start) {
            return turbo::invalid_argument_error(
                    "week_start must follow ISO convention (Monday=1, Sunday=7). Got week_start=",
                    options.week_start);
        }
        return turbo::OkStatus();
    }

    static inline turbo::Result<std::locale> GetLocale(const std::string &locale) {
        try {
            return std::locale(locale.c_str());
        } catch (const std::runtime_error &ex) {
            return turbo::invalid_argument_error("Cannot find locale '", locale, "': ", ex.what());
        }
    }

    struct NonZonedLocalizer {
        using days_t = sys_days;

        // No-op conversions: UTC -> UTC
        template<typename Duration>
        sys_time <Duration> ConvertTimePoint(int64_t t) const {
            return sys_time<Duration>(Duration{t});
        }

        template<typename Duration>
        Duration ConvertLocalToSys(Duration t, turbo::Status *st) const {
            return t;
        }

        sys_days ConvertDays(sys_days d) const { return d; }
    };

    struct ZonedLocalizer {
        using days_t = local_days;

        // Timezone-localizing conversions: UTC -> local time
        const time_zone *tz;

        template<typename Duration>
        local_time <Duration> ConvertTimePoint(int64_t t) const {
            return tz->to_local(sys_time<Duration>(Duration{t}));
        }

        template<typename Duration>
        Duration ConvertLocalToSys(Duration t, turbo::Status *st) const {
            try {
                return zoned_time<Duration>{tz, local_time<Duration>(t)}
                        .get_sys_time()
                        .time_since_epoch();
            } catch (const alkaid_vendored::date::nonexistent_local_time &e) {
                *st = turbo::invalid_argument_error("Local time does not exist: ", e.what());
                return Duration{0};
            } catch (const alkaid_vendored::date::ambiguous_local_time &e) {
                *st = turbo::invalid_argument_error("Local time is ambiguous: ", e.what());
                return Duration{0};
            }
        }

        local_days ConvertDays(sys_days d) const { return local_days(year_month_day(d)); }
    };

    template<typename Duration>
    struct TimestampFormatter {
        const char *format;
        const time_zone *tz;
        std::ostringstream bufstream;

        explicit TimestampFormatter(const std::string &format, const time_zone *tz,
                                    const std::locale &locale)
                : format(format.c_str()), tz(tz) {
            bufstream.imbue(locale);
            // Propagate errors as C++ exceptions (to get an actual error message)
            bufstream.exceptions(std::ios::failbit | std::ios::badbit);
        }

        turbo::Result<std::string> operator()(int64_t arg) {
            bufstream.str("");
            const auto zt = zoned_time<Duration>{tz, sys_time<Duration>(Duration{arg})};
            try {
                alkaid_vendored::date::to_stream(bufstream, format, zt);
            } catch (const std::runtime_error &ex) {
                bufstream.clear();
                return turbo::invalid_argument_error("Failed formatting timestamp: ", ex.what());
            }
            // XXX could return a view with std::ostringstream::view() (C++20)
            return std::move(bufstream).str();
        }
    };

//
// Which types to generate a kernel for
//
    struct WithDates {
    };
    struct WithTimes {
    };
    struct WithTimestamps {
    };
    struct WithStringTypes {
    };

// This helper allows generating temporal kernels for selected type categories
// without any spurious code generation for other categories (e.g. avoid
// generating code for date kernels for a times-only function).
    template<typename Factory>
    void add_temporal_kernels(Factory *fac) {}

    template<typename Factory, typename... WithOthers>
    void add_temporal_kernels(Factory *fac, WithDates, WithOthers... others) {
        fac->template add_kernel<days, Date32Type>(date32());
        fac->template add_kernel<std::chrono::milliseconds, Date64Type>(date64());
        add_temporal_kernels(fac, std::forward<WithOthers>(others)...);
    }

    template<typename Factory, typename... WithOthers>
    void add_temporal_kernels(Factory *fac, WithTimes, WithOthers... others) {
        fac->template add_kernel<std::chrono::seconds, Time32Type>(time32(TimeUnit::SECOND));
        fac->template add_kernel<std::chrono::milliseconds, Time32Type>(time32(TimeUnit::MILLI));
        fac->template add_kernel<std::chrono::microseconds, Time64Type>(time64(TimeUnit::MICRO));
        fac->template add_kernel<std::chrono::nanoseconds, Time64Type>(time64(TimeUnit::NANO));
        add_temporal_kernels(fac, std::forward<WithOthers>(others)...);
    }

    template<typename Factory, typename... WithOthers>
    void add_temporal_kernels(Factory *fac, WithTimestamps, WithOthers... others) {
        fac->template add_kernel<std::chrono::seconds, TimestampType>(
                match::TimestampTypeUnit(TimeUnit::SECOND));
        fac->template add_kernel<std::chrono::milliseconds, TimestampType>(
                match::TimestampTypeUnit(TimeUnit::MILLI));
        fac->template add_kernel<std::chrono::microseconds, TimestampType>(
                match::TimestampTypeUnit(TimeUnit::MICRO));
        fac->template add_kernel<std::chrono::nanoseconds, TimestampType>(
                match::TimestampTypeUnit(TimeUnit::NANO));
        add_temporal_kernels(fac, std::forward<WithOthers>(others)...);
    }

    template<typename Factory, typename... WithOthers>
    void add_temporal_kernels(Factory *fac, WithStringTypes, WithOthers... others) {
        fac->template add_kernel<TimestampType, StringType>(utf8());
        fac->template add_kernel<TimestampType, LargeStringType>(large_utf8());
        add_temporal_kernels(fac, std::forward<WithOthers>(others)...);
    }

    //
    // Executor class for temporal component extractors, i.e. scalar kernels
    // with the signature Timestamp -> <non-temporal scalar type `OutType`>
    //
    // The `Op` parameter is templated on the Duration (which depends on the timestamp
    // unit) and a Localizer class (depending on whether the timestamp has a
    // timezone defined).
    //
    template<template<typename...> class Op, typename Duration, typename InType,
            typename OutType, typename... Args>
    struct TemporalComponentExtractBase {
        template<typename OptionsType>
        static turbo::Status exec_with_options(KernelContext *ctx, const OptionsType *options,
                                             const ExecSpan &batch, ExecResult *out, Args... args) {
            const auto &timezone = GetInputTimezone(*batch[0].type());
            if (timezone.empty()) {
                using ExecTemplate = Op<Duration, NonZonedLocalizer>;
                auto op = ExecTemplate(options, NonZonedLocalizer(), args...);
                applicator::ScalarUnaryNotNullStateful<OutType, InType, ExecTemplate> kernel{op};
                return kernel.Exec(ctx, batch, out);
            } else {
                TURBO_MOVE_OR_RAISE(auto tz, LocateZone(timezone));
                using ExecTemplate = Op<Duration, ZonedLocalizer>;
                auto op = ExecTemplate(options, ZonedLocalizer{tz}, args...);
                applicator::ScalarUnaryNotNullStateful<OutType, InType, ExecTemplate> kernel{op};
                return kernel.Exec(ctx, batch, out);
            }
        }
    };

    template<template<typename...> class Op, typename OutType>
    struct TemporalComponentExtractBase<Op, days, Date32Type, OutType> {
        template<typename OptionsType>
        static turbo::Status exec_with_options(KernelContext *ctx, const OptionsType *options,
                                             const ExecSpan &batch, ExecResult *out) {
            using ExecTemplate = Op<days, NonZonedLocalizer>;
            auto op = ExecTemplate(options, NonZonedLocalizer());
            applicator::ScalarUnaryNotNullStateful<OutType, Date32Type, ExecTemplate> kernel{op};
            return kernel.Exec(ctx, batch, out);
        }
    };

    template<template<typename...> class Op, typename OutType>
    struct TemporalComponentExtractBase<Op, std::chrono::milliseconds, Date64Type, OutType> {
        template<typename OptionsType>
        static turbo::Status exec_with_options(KernelContext *ctx, const OptionsType *options,
                                             const ExecSpan &batch, ExecResult *out) {
            using ExecTemplate = Op<std::chrono::milliseconds, NonZonedLocalizer>;
            auto op = ExecTemplate(options, NonZonedLocalizer());
            applicator::ScalarUnaryNotNullStateful<OutType, Date64Type, ExecTemplate> kernel{op};
            return kernel.Exec(ctx, batch, out);
        }
    };

    template<template<typename...> class Op, typename OutType>
    struct TemporalComponentExtractBase<Op, std::chrono::seconds, Time32Type, OutType> {
        template<typename OptionsType>
        static turbo::Status exec_with_options(KernelContext *ctx, const OptionsType *options,
                                             const ExecSpan &batch, ExecResult *out) {
            using ExecTemplate = Op<std::chrono::seconds, NonZonedLocalizer>;
            auto op = ExecTemplate(options, NonZonedLocalizer());
            applicator::ScalarUnaryNotNullStateful<OutType, Time32Type, ExecTemplate> kernel{op};
            return kernel.Exec(ctx, batch, out);
        }
    };

    template<template<typename...> class Op, typename OutType>
    struct TemporalComponentExtractBase<Op, std::chrono::milliseconds, Time32Type, OutType> {
        template<typename OptionsType>
        static turbo::Status exec_with_options(KernelContext *ctx, const OptionsType *options,
                                             const ExecSpan &batch, ExecResult *out) {
            using ExecTemplate = Op<std::chrono::milliseconds, NonZonedLocalizer>;
            auto op = ExecTemplate(options, NonZonedLocalizer());
            applicator::ScalarUnaryNotNullStateful<OutType, Time32Type, ExecTemplate> kernel{op};
            return kernel.Exec(ctx, batch, out);
        }
    };

    template<template<typename...> class Op, typename OutType>
    struct TemporalComponentExtractBase<Op, std::chrono::microseconds, Time64Type, OutType> {
        template<typename OptionsType>
        static turbo::Status exec_with_options(KernelContext *ctx, const OptionsType *options,
                                             const ExecSpan &batch, ExecResult *out) {
            using ExecTemplate = Op<std::chrono::microseconds, NonZonedLocalizer>;
            auto op = ExecTemplate(options, NonZonedLocalizer());
            applicator::ScalarUnaryNotNullStateful<OutType, Date64Type, ExecTemplate> kernel{op};
            return kernel.Exec(ctx, batch, out);
        }
    };

    template<template<typename...> class Op, typename OutType>
    struct TemporalComponentExtractBase<Op, std::chrono::nanoseconds, Time64Type, OutType> {
        template<typename OptionsType>
        static turbo::Status exec_with_options(KernelContext *ctx, const OptionsType *options,
                                             const ExecSpan &batch, ExecResult *out) {
            using ExecTemplate = Op<std::chrono::nanoseconds, NonZonedLocalizer>;
            auto op = ExecTemplate(options, NonZonedLocalizer());
            applicator::ScalarUnaryNotNullStateful<OutType, Date64Type, ExecTemplate> kernel{op};
            return kernel.Exec(ctx, batch, out);
        }
    };

    template<template<typename...> class Op, typename Duration, typename InType,
            typename OutType, typename... Args>
    struct TemporalComponentExtract
            : public TemporalComponentExtractBase<Op, Duration, InType, OutType, Args...> {
        using Base = TemporalComponentExtractBase<Op, Duration, InType, OutType, Args...>;

        static turbo::Status Exec(KernelContext *ctx, const ExecSpan &batch, ExecResult *out,
                                  Args... args) {
            const FunctionOptions *options = nullptr;
            return Base::exec_with_options(ctx, options, batch, out, args...);
        }
    };

}  // namespace nebula::compute::internal
